package skewreduce.framework.planner;

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.Console;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.PriorityQueue;
import java.util.SortedSet;
import java.util.TreeSet;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileStatus;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapred.ClusterStatus;
import org.apache.hadoop.mapred.JobClient;
import org.apache.hadoop.mapred.JobConf;
import org.apache.hadoop.util.GenericOptionsParser;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import skewreduce.framework.Schedule;
import skewreduce.framework.logical.LogicalOp;
import skewreduce.framework.physical.PhysicalOp;
import skewreduce.framework.planner.partition.Partitioner;
import skewreduce.lib.ITuple;
import skewreduce.lib.KahanSum;
import skewreduce.lib.Partition;


/**
 * using sample and known execution, fit constants for different cost model.
 * @author yongchul
 *
 */
public class Shell {
	private static final Logger LOG = LoggerFactory.getLogger(Shell.class);
	
	final Configuration conf;
	final FileSystem fs;
	List<ITuple> sample = new ArrayList<ITuple>();
	
	private IdentityHashMap<PartitionNode,PhysicalOp.ID> node2id = new IdentityHashMap<PartitionNode,PhysicalOp.ID>();
	private HashMap<PhysicalOp.ID,PartitionNode> id2node = new HashMap<PhysicalOp.ID,PartitionNode>();

	private HashMap<PhysicalOp.ID,Double> id2time = new HashMap<PhysicalOp.ID,Double>();
	private HashSet<PartitionNode> currentPlan = new HashSet<PartitionNode>(); // for scheduler
	private SortedSet<PartitionNode> pending = new TreeSet<PartitionNode>();
	
	class LinearRegression {
		ArrayList<Double> x; 
		ArrayList<Double> y;

		KahanSum sumx2 = new KahanSum();
		KahanSum sumx = new KahanSum();
		KahanSum sumy = new KahanSum();
		
		KahanSum xxbar = new KahanSum();
		KahanSum yybar = new KahanSum();
		KahanSum xybar = new KahanSum();
		
		double xbar;
		double ybar;
		
		double beta0;
		double beta1;
		
        KahanSum rss = new KahanSum();      // residual sum of squares
        KahanSum ssr = new KahanSum();      // regression sum of squares
		
		LinearRegression() {
			this(10);
		}
		
		LinearRegression(int cap) {
			x = new ArrayList<Double>(cap);
			y = new ArrayList<Double>(cap);
		}
		
		public void reset() {
			x.clear(); y.clear();
			sumx.reset();
			sumy.reset();
			sumx2.reset();
			xxbar.reset();
			yybar.reset();
			xybar.reset();
		}
		public void update(double x,double y) {
			this.x.add(x);
			this.y.add(y);
			
			sumx.add(x);
			sumy.add(y);
			sumx2.add(x*x);
		}
		
		public void doFinal() {
			xbar = sumx.value() / x.size();
			ybar = sumy.value() / x.size();
			
			for ( int i = 0; i < x.size(); ++i ) {
				double xx = x.get(i);
				double yy = y.get(i);
				
				xxbar.add((xx - xbar)*(xx - xbar));
				yybar.add((yy - ybar)*(yy - ybar));
				xybar.add((xx - xbar)*(yy - ybar));
			}
			
			beta1 = xybar.value() / xxbar.value();
			beta0 = ybar - beta1 * xbar;
			
	        for (int i = 0; i < x.size(); i++) {
	            double fit = beta1*x.get(i) + beta0;
	            rss.add((fit - y.get(i)) * (fit - y.get(i)));
	            ssr.add((fit - ybar) * (fit - ybar));
	        }
		}
		
		public double getBeta0() { return beta0; }
		public double getBeta1() { return beta1; }
		public double getR2() { return ssr.value() / yybar.value(); }
				
		@Override
		public String toString() {
			return String.format("y = %g * x + %g (R^2 = %g)",beta1,beta0,getR2());
		}

        // analyze results
//        int df = n - 2;
//        double svar  = rss / df;
//        double svar1 = svar / xxbar;
//        double svar0 = svar/n + xbar*xbar*svar1;
//        System.out.println("R^2                 = " + R2);
//        System.out.println("std error of beta_1 = " + Math.sqrt(svar1));
//        System.out.println("std error of beta_0 = " + Math.sqrt(svar0));
//        svar0 = svar * sumx2 / (n * xxbar);
//        System.out.println("std error of beta_0 = " + Math.sqrt(svar0));
//
//        System.out.println("SSTO = " + yybar);
//        System.out.println("SSE  = " + rss);
//        System.out.println("SSR  = " + ssr);
	}

	private EaggCostModel model;
	private PartitionNode rootNode;
	private Partition rootPartition;
	private double sampleRate;
	private double bandWidth;
	
	private double costWorkPerSec;
	private double costMergePerSec;

	private Partitioner partitioner;
	private PriorityScheduler scheduler;
	
	private static final int DEFAULT_SAMPLE_SIZE = 10;
	private int sampleWorkSize = DEFAULT_SAMPLE_SIZE;
	private int sampleMergeSize = DEFAULT_SAMPLE_SIZE;

	LinearRegression regWork = new LinearRegression();
	LinearRegression regMerge = new LinearRegression();
	
	private boolean verbose;
	
	private void partition(PartitionNode node) {
		ArrayList<ITuple> ls = new ArrayList<ITuple>();
		ArrayList<ITuple> rs = new ArrayList<ITuple>();
		
		Partition parent = node.getPartition();
		
		for ( ITuple t : node.getSample() ) {
			int child = parent.whichChild(t);
			if ( child == 0 ) ls.add(t);
			else rs.add(t);
		}
		
		ls.trimToSize();
		rs.trimToSize();
		
		node.setChildren(new PartitionNode(node,parent.getLeft(),ls),new PartitionNode(node,parent.getRight(),rs));
	}
	
	private void buildPartitionTree(PartitionNode node) {
		if ( node == null ) return;
		
		if ( ! node.getPartition().isLeaf() ) {
			partition(node);
			
			buildPartitionTree(node.getLeft());
			buildPartitionTree(node.getRight());
		}

		PhysicalOp.ID id = new PhysicalOp.ID(node.getPartition());
		node2id.put(node,id);
		id2node.put(id,node);
		currentPlan.add(node);
	}
	
	private void fit(PartitionNode node) {
		if ( node.isLeaf() ) {
			// estimate cost
			double cost = model.costWork(node);
			double time = id2time.get( node2id.get(node) );
			
			// fit const
			if ( sampleWorkSize-- > 0 ) {
				regWork.update(cost,time);
			}
		} else {
			fit(node.getLeft());
			fit(node.getRight());
			
			double cost = model.costMerge(node,node.getLeft(),node.getRight());
			double time = id2time.get(node2id.get(node));
			
			// fit cost
			if ( sampleMergeSize-- > 0 )
				regMerge.update(cost, time);
		}
	}
	
	private void fit() {
		fit(rootNode);
		
		regWork.doFinal();
		regMerge.doFinal();
	}
	
	private void buildPartitionTree() {
		rootNode = new PartitionNode(rootPartition,sample);
		buildPartitionTree(rootNode);
	}
	
	
	private PartitionNode lookup(String s) {
		char c = s.charAt(0);
		if ( '0' <= c && c <= '9' ) {
			s = "WORK-".concat(s);
		}
		PhysicalOp.ID id = new PhysicalOp.ID(s);
		PartitionNode node = id2node.get(id);
		if ( node == null ) {
			if ( id.getType() == LogicalOp.Type.MERGE ) {
				id = new PhysicalOp.ID("WORK",id);
			} else {
				id = new PhysicalOp.ID("MERGE",id);
			}
			node = id2node.get(id);
		}
		return node;
	}
	
	//////////////////////////////
	
	Shell(Configuration cnf) throws IOException {
		conf = cnf;
		fs = FileSystem.get(this.conf);
	}
	
	private int loadSample(Path fn) throws IOException {
		int before = sample.size();
		FSDataInputStream in = null;
		ITuple t = null;
		try {
			in = fs.open(fn);
			while ( (t = model.createTuple(in)) != null ) {
				sample.add(t);
			}
		} catch ( EOFException ignore ) {
			// we simply ignore this exception and throw all other IOExceptions
		} finally {
			if ( in != null ) try { in.close(); } catch ( IOException ignore ) {}
			in = null;
		}
		return sample.size() - before;
	}
	
	private void doRollback(Console console) {
		// undo the last partition
		if ( pivot != null ) pivot.rollback();
		if ( ! pending.isEmpty() ) {
			console.printf("discard pending %d changes\n", pending.size());
		}
		pending.clear();
	}
	
	private void doCommit(Console console) {
		currentPlan.addAll(pending);
		for ( PartitionNode node : pending ) {
			id2node.put(node.getOpID(), node);
		}
		pending.clear();
	}
	
	private boolean confirmCommit(Console console) {
		if ( ! pending.isEmpty() ) {
			while ( true ) {
				String line = console.readLine("There are pending changes. Want to commit first? [yes|no]");
				if ( line == null )
					return false;
				
				line = line.trim();
				if ( "yes".equalsIgnoreCase(line) ) {
					doCommit(console);
					break;
				} else if ( "no".equalsIgnoreCase(line) ) {
					doRollback(console);
					break;
				}
			}
		}
		return true;
	}
	
	private boolean doPartitionUntil(Console console) {
		if ( pivot == null ) {
			System.err.println("usage: partition [node-id]");
		} else if ( ! pivot.isLeaf() ) {
			System.err.printf("can't partition %s because it is non-leaf node.\n",pivot.getOpID());
		} else if ( ! pivot.canPartition() ) {
			System.err.printf("can't repartition %s because there is too few sample.\n",pivot.getOpID());
		} else {
			// do repartition
			double costwork = pivot.getWorkCost();
			if ( costwork == 0.0 || Double.isNaN(costwork) ) {
				costwork = model.costWork(pivot);
				pivot.setWorkCost(costwork, costWorkPerSec);
			}
			double bestCost = pivot.getWorkSecs();

			PriorityQueue<PartitionNode> queue = new PriorityQueue<PartitionNode>();
			queue.add(pivot);
			pending.add(pivot);
			
			while ( ! queue.isEmpty() ) {
				PartitionNode node = queue.poll();
				
				PartitionNode[] children = partitioner.partition(node);
				if ( children == null ) {
					continue;
				} else {
					PartitionNode l = children[0];
					PartitionNode r = children[1];
					
					if ( LOG.isDebugEnabled() ) {
						LOG.debug("expected items = {} / max items = {}",node.getExpectedItems(sampleRate), model.getMaxWorkDataItems() );
					}
					
					console.printf("Total min time = %g, Previous time = %g\n", node.getTotalMinSecs(model.getScheduleDelay()), node.getWorkSecs());
					console.printf("Merge cost = %g/time = %g\n", node.getCost(), node.getMergeSecs());
					console.printf("Left = %s -- cost = %g/time = %g\n", l.getPartition(), l.getCost(), l.getWorkSecs());
					console.printf("Right = %s -- cost = %g/time = %g\n", r.getPartition(), r.getCost(), r.getWorkSecs());
					
					if ( node.getTotalMinSecs(model.getScheduleDelay()) < node.getWorkSecs() ) {
						scheduler.reset(pending);
						scheduler.add(l);
						scheduler.add(r);
						double schedCost = scheduler.estimateSchedule();
						
						console.printf("schedule cost = %g/best cost = %g\n",schedCost,bestCost);

						if ( LOG.isDebugEnabled() ) {
							LOG.debug("best cost = {}, current = {}",bestCost,schedCost);
						}
						
						if ( schedCost < bestCost ) {
							bestCost = schedCost;
							if ( LOG.isDebugEnabled() ) {
								LOG.debug("partition {}, total cost = {}",node.partition.toSpec(), bestCost);
							}
							
							// add two split nodes to leaves and insert to the queue
							pending.add(l); pending.add(r);
							queue.add(l); queue.add(r);
							
							console.printf("KEEP PARTITION\n");
						} else {
							console.printf("DISCARD PARTITION - no schedule benefit\n");
							node.rollback(); // drop children
						}
					} else {
						console.printf("DISCARD PARTITION - no benefit\n");
						node.rollback(); // drop children
					}
				}
			}
		}
		
		return this.confirmCommit(console);
	}
	
	private PartitionNode pivot;
	
	private boolean handleCommand(Console console,String cmd, String[] args) throws Exception {
		
		if ( cmd.startsWith("quit") || cmd.startsWith("exit") ) {
			return false;
		} else if ( "select".equalsIgnoreCase(cmd) ) {
			if ( args.length == 0 ) {
				System.err.println("usage: select node-id");
				pivot = null;
			} else if ( args[0].equalsIgnoreCase("none") ) {
				console.printf("selected none\n");
				pivot = null;
			} else {
				PartitionNode node = lookup(args[0]);
				if ( node == null ) {
					System.err.println("Could not locate given node: "+args[0]);
					pivot = null;
				} else {
					pivot = node;
					console.printf("node %s was selected\n", node.getOpID());
				}
			}
		} else if ( cmd.equalsIgnoreCase("partition") ) {
			if ( args.length > 0 ) {
				pivot = lookup(args[0]);
			}
			if ( pivot == null ) {
				System.err.println("usage: partition [node-id]");
			} else if ( ! pivot.isLeaf() ) {
				System.err.printf("can't partition %s because it is non-leaf node.\n",pivot.getOpID());
			} else if ( ! pivot.canPartition() ) {
				System.err.printf("can't repartition %s because there is too few sample.\n",pivot.getOpID());
			} else {
				// do repartition
				if ( LOG.isDebugEnabled() ) {
					LOG.debug("Considering partitioning {} - cost {}",pivot.getPartition(),pivot.getCost());
				}

				PartitionNode[] children = partitioner.partition(pivot);
				if ( children == null ) {
					System.err.printf("can't find a feasible partitioning plan for %s.\n", pivot.getOpID());
				} else {
					PartitionNode l = children[0];
					PartitionNode r = children[1];
					
					if ( LOG.isDebugEnabled() ) {
						LOG.debug("expected items = {} / max items = {}",pivot.getExpectedItems(sampleRate), model.getMaxWorkDataItems() );
					}
					
					pending.add(l);
					pending.add(r);

					// FIXME: print out the statistics
					console.printf("Merge cost = %g/time = %g\n", pivot.getCost(), pivot.getMergeSecs());
					console.printf("Left = %s -- cost = %g/time = %g\n", l.getPartition(), l.getCost(), l.getWorkSecs());
					console.printf("Right = %s -- cost = %g/time = %g\n", r.getPartition(), r.getCost(), r.getWorkSecs());
				}	
			}
		} else if ( cmd.equalsIgnoreCase("partitionUntil") ) {
			// repartition until there is no further improvement assuming given number of machines
			int argc = 0;
			boolean bad = false;
			int numMachines = 1;
			for ( argc = 0; argc < args.length; ++argc ) {
				if ( args[argc].charAt(0) != '-' ) break;
				if ( "-n".equals(args[argc]) ) {
					numMachines = Integer.parseInt(args[++argc]);
				} else {
					System.err.println("Unknown argument: "+args[argc]);
					bad = true;
					break;
				}
			}
			if ( bad ) {
				return true;
			}
			
			scheduler.setNumSlots(numMachines);
			
			if ( argc < args.length ) {
				for ( ; argc < args.length; ++argc ) {
					pivot = lookup(args[argc]);
					if ( pivot == null ) {
						System.err.printf("node %s is not found\n",args[argc]);
					} else {
						doPartitionUntil(console);
					}
				}
			} else {
				doPartitionUntil(console); // using current pivot
			}
		} else if ( cmd.startsWith("estimate") ) {
			int numMachines = 1;
			if ( args.length > 0 ) {
				numMachines = Integer.parseInt(args[0]);
			}
			scheduler.setNumSlots(numMachines);
			scheduler.reset(currentPlan);
			double t = scheduler.estimateSchedule();
			console.printf("Estimated time = %g seconds using %d machines\n", t,numMachines);
		} else if ( cmd.startsWith("print") ) {
			if ( pivot == null ) {
				rootPartition.print(System.out);
			} else {
				pivot.partition.print(System.out);
			}
		} else if ( cmd.startsWith("store") || cmd.startsWith("save") ) {
			// save the current partition and schedule
			confirmCommit(console);
			doSave(console,args);
		} else if ( cmd.equalsIgnoreCase("commit") ) {
			doCommit(console);
		} else if ( cmd.equalsIgnoreCase("rollback") ) {
			doRollback(console);
		} else if ( cmd.equalsIgnoreCase("set") ) {
			if ( args.length == 0 ) {
				System.err.println("usage: set varname value");
			} else {
				if ( args.length > 1 ) {
					conf.set(args[0], args[1]);
				}
				console.printf("%s = %s\n", args[0], conf.get(args[0]) );
			}
		} else {
			System.err.println("Unknown command: "+cmd);
		}
		
		console.flush();
		
		return true;
	}
	
	private void doSave(Console console, String[] args) throws IOException {
		String fn = args[0];
		FileOutputStream fos = null;
		DataOutputStream out = null;
		try {
			fos = new FileOutputStream(String.format("%s.partition",fn));
			out = new DataOutputStream(new BufferedOutputStream(fos,4096*1024));
			rootPartition.write(out);
		} finally {
			if ( out != null ) try { out.close(); } catch ( IOException ignore ) {}
			out = null;
			if ( fos != null ) try { fos.close(); } catch ( IOException ignore ) {}
			fos = null;
		}
		
		// sanity check
		for ( PartitionNode node : currentPlan ) {
			if ( node.isLeaf() ) {
				if ( Double.isNaN(node.getWorkCost()) || node.getWorkCost() == 0.0 ) {
					node.setWorkCost(model.costWork(node), costWorkPerSec);
				}
			} else {
				if ( Double.isNaN(node.getMergeCost()) ) {
					node.setMergeCost(model.costMerge(node,node.getLeft(),node.getRight()), costMergePerSec);
				}
			}
		}
		
		scheduler.setNumSlots(1); // longest job first
		scheduler.reset(currentPlan);
		scheduler.estimateSchedule();
		
		PrintStream scheduleOut = null;
		try {
			scheduleOut = new PrintStream(String.format("%s.schedule",fn));
			for ( Schedule s : scheduler.getSchedule() ) {
				scheduleOut.println(s);
			}
		} finally {
			if ( scheduleOut != null ) scheduleOut.close();
			scheduleOut = null;
		}
	}

	private void doMain(String[] args) throws Exception {
		boolean bad = false;
		String planFile = null;
		String timeFile = null;
		boolean printAndQuit = false;
		int i = 0;
		
		for ( ; i < args.length; ++i ) {
			if ( args[i].charAt(0) != '-' )
				break;
			
			if ( "-sampleRate".equals(args[i]) ) {
				sampleRate = Double.parseDouble(args[++i]);
				conf.setFloat("skewreduce.planner.sampleRate", (float)sampleRate);
			} else if ( "-model".equals(args[i]) ) {
				conf.set(EaggCostModel.MODEL_CLASS_PROP, args[++i]);
			} else if ( "-bandWidth".equals(args[i]) ) {
				bandWidth = Double.parseDouble(args[++i]);
				conf.setFloat("skewreduce.planner.partition.bandWidth", (float)bandWidth);
			} else if ( "-verbose".equals(args[i]) || "-v".equals(args[i]) ) {
				verbose = true;
			} else if ( "-plan".equals(args[i]) ) {
				planFile = args[++i];
			} else if ( "-time".equals(args[i]) ) {
				timeFile = args[++i];
			} else if ( "-sampleSize".equals(args[i]) ) {
				int sz = Integer.parseInt(args[++i]);
				this.sampleWorkSize = this.sampleMergeSize = sz;
			} else if ( "-w2s".equals(args[i]) ) {
				costWorkPerSec = Double.parseDouble(args[++i]);
				conf.setFloat("skewreduce.planner.cost.work2sec", (float)costWorkPerSec);
			} else if ( "-m2s".equals(args[i]) ) {
				costMergePerSec = Double.parseDouble(args[++i]);
				conf.setFloat("skewreduce.planner.cost.merge2sec",(float)costMergePerSec);
			} else if ( "-delay".equals(args[i]) ) {
				float delay = Float.parseFloat(args[++i]);
				conf.setFloat("skewreduce.planner.scheduler.delay",delay);
			} else if ( "-print".equals(args[i]) ) {
				printAndQuit = true;
			} else {
				System.err.println("Unknown option: "+args[i]);
				bad = true;
			}
		}
		
		model = EaggCostModel.getInstance(conf);
		
		bad = bad || model == null;

		if ( ! bad ) {
			for ( ; i < args.length; ++i ) {
				FileStatus[] fns = fs.globStatus(new Path(args[i]));
				for ( FileStatus fn : fns ) {
					int nsample = loadSample(fn.getPath());
					LOG.info("{} samples have loaded from {}",nsample,fn.getPath());
	
				}
			}
		}
		
		LOG.info("Total {} samples have been loaded",sample.size());
		
		if ( bad || ( sample.isEmpty() && ! printAndQuit ) ) {
			// FIXME print usage information
			System.exit(-1);
		}
		
		rootPartition = model.getRootPartition();
		{
			// load plan
			DataInputStream input = new DataInputStream(new FileInputStream(planFile));
			rootPartition.readFields(input);
			input.close();
		}
		
		if ( timeFile != null ) {
			// load timing
			BufferedReader input = new BufferedReader(new FileReader(timeFile));
			String line;
			while ( (line = input.readLine()) != null ) {
				String[] fields = line.split(",");
				PhysicalOp.ID id = new PhysicalOp.ID(fields[0]);
				if ( id.getType() != LogicalOp.Type.WORK && id.getType() != LogicalOp.Type.MERGE ) continue;
				double t = Double.parseDouble(fields[1]);
				id2time.put(id, t);
			}
			input.close();
		}
		
		buildPartitionTree();
		
		if ( timeFile != null ) {
			fit(); // bottom-up evaluation
			
			System.out.println("Work Estimate: "+ this.regWork);
			System.out.println("Merge Estimate: "+ this.regMerge);
		} else {
			partitioner = Partitioner.getInstance(model,null); // default brute-force
			scheduler = new PriorityScheduler(conf);

			// get input from console and run command
			Console console = System.console();
			String line = null;
			
			if (printAndQuit) {
				handleCommand(console,"print",new String[0]);
				handleCommand(console,"quit",new String[0]);
			} else {
				while ( (line = console.readLine("EaggShell>")) != null ) {
					line = line.trim();
					if ( line.length() == 0 ) continue;
					
					String[] xxx = line.split("\\s+");
					String cmd = xxx[0];
					String[] params = Arrays.copyOfRange(xxx, 1, xxx.length);
					try {
						if ( ! handleCommand(console,cmd,params) ) {
							break;
						}
					} catch ( Exception ex ) {
						System.err.println("Exception caught: "+ex.getMessage());
						if ( verbose ) {
							ex.printStackTrace();
						}
					}
				}
			}
		}
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) throws Exception {
        GenericOptionsParser options = new GenericOptionsParser(new Configuration(),args);
        Configuration conf = options.getConfiguration();
        String[] rest = options.getRemainingArgs();
        
        //LOG.info("current working directory = {}",System.getProperty("user.dir"));

        if ( rest == null || rest.length == 0 ) {
            System.err.println("No class argument is given!");
            System.exit(1);
        }
        
        Shell prog = new Shell(conf);
        prog.doMain(rest);		
	}

}
